import logging
import yaml
import os
import copy
import argparse
from pathlib import Path
from typing import Dict, Any, Optional, Union, List, Tuple, Set
import pandas as pd
import concurrent.futures
import threading

# 添加项目根目录到路径
import sys

sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), ".."))

from data_processor.loaders.base_loader import BaseDataLoader, AnomalyRecord
from data_processor.loaders.loader_factory import LoaderFactory

# 配置日志
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
# 创建一个锁用于多线程日志同步
log_lock = threading.Lock()


class DataPreparation:
    """
    数据准备类，负责加载和准备数据，支持多线程处理

    主要功能:
    1. 加载数据源配置
    2. 根据数据集名称（路径最后一部分）组织数据
    3. 为每个数据集创建对应的输出目录
    4. 加载每个数据集的异常记录
    5. 支持多线程并行处理数据
    """

    def __init__(self, config: Dict[str, Any], output_base_dir: Union[str, Path]):
        """
        初始化数据准备工作流

        参数:
            config: 配置字典，包含数据源、处理参数等
            output_base_dir: 输出基础目录
        """
        logger.info("开始初始化数据准备工作流")
        self.config = config
        self.exp_id = config.get("exp_id", "default")
        self.output_base_dir = Path(output_base_dir) / f"run_{self.exp_id}"
        self.output_base_dir.mkdir(parents=True, exist_ok=True)

        # ===== 数据处理配置 =====
        self.data_processing_config = config.get("data_processing", {})

        # 设置工作线程数量，默认为系统CPU核心数的一半
        self.num_workers = self.data_processing_config.get("num_workers", os.cpu_count() // 2 or 1)
        # 确保num_workers至少为1
        self.num_workers = max(1, self.num_workers)
        logger.info(f"工作线程数量设置为: {self.num_workers}")

        # 组织数据源按数据集名称（目录最后一部分）
        self.organized_data_sources = self._organize_data_sources_by_name()
        logger.info(f"数据准备工作流初始化完成，实验ID: {self.exp_id}")

    def _organize_data_sources_by_name(self) -> Dict[str, List[Dict[str, Any]]]:
        """
        将数据源按数据集名称（路径最后一部分）组织

        返回:
            Dict[str, List[Dict[str, Any]]]: 按数据集名称组织的数据源字典
        """
        organized_sources = {}
        data_sources = self.config.get("data_sources", [])

        for ds_config in data_sources:
            base_dir = ds_config.get("base_dir", "")
            if not base_dir:
                logger.warning(f"数据源配置缺少base_dir字段，跳过: {ds_config}")
                continue

            # 提取数据集名称（目录路径最后一部分）
            dataset_name = Path(base_dir).name

            # 将数据源添加到对应数据集名称的列表中
            if dataset_name not in organized_sources:
                organized_sources[dataset_name] = []
            organized_sources[dataset_name].append(ds_config)

            logger.info(f"数据集 '{dataset_name}' 添加了数据源: {base_dir}")

        logger.info(f"共组织了 {len(organized_sources)} 个数据集")
        return organized_sources

    def _load_datasource_anomaly_records(
        self, ds_config: Dict[str, Any], dataset_name: str
    ) -> Tuple[List[AnomalyRecord], Optional[BaseDataLoader]]:
        """
        加载单个数据源的异常记录

        参数:
            ds_config: 数据源配置字典
            dataset_name: 数据集名称

        返回:
            Tuple[List[AnomalyRecord], Optional[BaseDataLoader]]:
            (异常记录列表, 加载器实例)
        """
        ds_type = ds_config.get("type", "").lower()
        base_dir = ds_config.get("base_dir")

        # 创建以数据集名称命名的输出目录
        ds_output_dir = self.output_base_dir / dataset_name
        ds_output_dir.mkdir(parents=True, exist_ok=True)

        loader = LoaderFactory.create_loader(
            ds_config, ds_output_dir, self.data_processing_config
        )
        if not loader:
            with log_lock:
                logger.error(
                    f"无法为数据集 {dataset_name} (类型: {ds_type}) 创建加载器，跳过"
                )
            return [], None

        with log_lock:
            logger.info(f"开始加载数据集 {dataset_name} (类型: {ds_type}) 的异常记录")
        anomaly_records = loader._load_anomaly_records()

        if not anomaly_records:
            with log_lock:
                logger.warning(
                    f"数据集 {dataset_name} (类型: {ds_type}) 未找到任何异常记录"
                )
            return [], loader

        with log_lock:
            logger.info(
                f"数据集 {dataset_name} (类型: {ds_type}) 成功加载 {len(anomaly_records)} 个异常记录"
            )
        return anomaly_records, loader

    def _process_single_record(
        self, record: AnomalyRecord, dataset_name: str, loader: BaseDataLoader
    ) -> Tuple[str, Any]:
        """
        处理单个异常记录，用于多线程并行处理

        参数:
            record: 异常记录
            dataset_name: 数据集名称
            loader: 数据加载器实例

        返回:
            Tuple[str, Any]: (异常ID, 实验数据)
        """
        anomaly_id = record.anomaly_id
        with log_lock:
            logger.info(f"开始加载数据集 {dataset_name} 的异常 {anomaly_id} 的实验数据")

        try:
            experiment_data = loader.load_experiment_data_from_record(record)
            if not experiment_data:
                with log_lock:
                    logger.error(
                        f"无法为数据集 {dataset_name} 的异常 {anomaly_id} 加载实验数据"
                    )
                return anomaly_id, None

            with log_lock:
                logger.info(f"数据集 {dataset_name} 的异常 {anomaly_id} 实验数据加载完成")
            return anomaly_id, experiment_data
        except Exception as e:
            with log_lock:
                logger.error(f"处理异常 {anomaly_id} 时出错: {e}")
            return anomaly_id, None

    def load_dataset_experiment_data(
        self,
        dataset_name: str,
        anomaly_records: List[AnomalyRecord],
        loader: BaseDataLoader,
    ) -> Dict[str, Any]:
        """
        加载单个数据集的所有实验数据，使用多线程并行处理

        参数:
            dataset_name: 数据集名称
            anomaly_records: 异常记录列表
            loader: 数据加载器实例

        返回:
            Dict[str, Any]: 实验数据字典，键为异常ID
        """
        experiment_data_dict = {}

        # 使用线程池并行处理异常记录
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.num_workers) as executor:
            # 提交所有任务
            future_to_record = {
                executor.submit(self._process_single_record, record, dataset_name, loader): record
                for record in anomaly_records
            }

            # 获取结果
            for future in concurrent.futures.as_completed(future_to_record):
                anomaly_id, experiment_data = future.result()
                if experiment_data:
                    experiment_data_dict[anomaly_id] = experiment_data

        with log_lock:
            logger.info(
                f"数据集 {dataset_name} 共加载了 {len(experiment_data_dict)} 个实验数据"
            )

        # 处理details.json，确保所有实验数据具有一致的实体和指标集
        if experiment_data_dict:
            with log_lock:
                logger.info(f"处理数据集 {dataset_name} 的details.json文件")
            details = loader.handle_details_json(experiment_data_dict)
            with log_lock:
                logger.info(f"数据集 {dataset_name} 的details.json处理完成")

        return experiment_data_dict

    def _process_dataset(self, dataset_name: str, data_sources: List[Dict[str, Any]]) -> Tuple[str, Dict[str, Any]]:
        """
        处理单个数据集，用于多线程并行处理

        参数:
            dataset_name: 数据集名称
            data_sources: 数据源配置列表

        返回:
            Tuple[str, Dict[str, Any]]: (数据集名称, 实验数据字典)
        """
        with log_lock:
            logger.info(f"开始处理数据集: {dataset_name}")

        # 对单个数据集的所有数据源进行处理
        all_records = []
        current_loader = None

        for ds_config in data_sources:
            records, loader = self._load_datasource_anomaly_records(
                ds_config, dataset_name
            )
            if records and loader:
                all_records.extend(records)
                current_loader = loader  # 使用最后一个成功的加载器

        if not all_records or not current_loader:
            with log_lock:
                logger.warning(
                    f"数据集 {dataset_name} 没有有效的异常记录或加载器，跳过"
                )
            return dataset_name, {}

        # 加载实验数据
        dataset_experiment_data = self.load_dataset_experiment_data(
            dataset_name, all_records, current_loader
        )

        return dataset_name, dataset_experiment_data

    def prepare_all_datasets(self) -> Dict[str, Dict[str, Any]]:
        """
        准备所有数据集的数据 - 主入口方法，使用多线程并行处理

        返回:
            Dict[str, Dict[str, Any]]: 按数据集名称组织的实验数据字典
        """
        logger.info("开始准备所有数据集")
        all_experiment_data = {}

        # 如果数据集数量少于线程数，调整线程数
        effective_workers = min(self.num_workers, len(self.organized_data_sources))

        # 使用线程池并行处理数据集
        with concurrent.futures.ThreadPoolExecutor(max_workers=effective_workers) as executor:
            # 提交所有任务
            future_to_dataset = {
                executor.submit(self._process_dataset, dataset_name, data_sources): dataset_name
                for dataset_name, data_sources in self.organized_data_sources.items()
            }

            # 获取结果
            for future in concurrent.futures.as_completed(future_to_dataset):
                dataset_name, dataset_experiment_data = future.result()
                if dataset_experiment_data:
                    all_experiment_data[dataset_name] = dataset_experiment_data

        logger.info(f"所有数据集准备完成，共 {len(all_experiment_data)} 个数据集")
        return all_experiment_data


def create_data_preparation(
    config_path: Union[str, Path], output_base_dir: Union[str, Path]
) -> Optional[DataPreparation]:
    """从配置文件创建数据准备实例"""
    config_path = Path(config_path)
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
    except Exception as e:
        logger.error(f"无法读取配置文件 {config_path}: {e}")
        return None

    if not config:
        logger.error(f"配置文件 {config_path} 为空或格式无效")
        return None

    return DataPreparation(config=config, output_base_dir=Path(output_base_dir))


# 准备数据并保存主函数
def prepare_and_save_data(
    config_path: Union[str, Path], output_base_dir: Union[str, Path]
) -> Dict[str, Dict[str, Any]]:
    """
    准备并保存数据的主函数

    参数:
        config_path: 配置文件路径
        output_base_dir: 输出基础目录

    返回:
        Dict[str, Dict[str, Any]]: 按数据集名称组织的实验数据字典
    """
    # 创建数据准备实例
    data_preparation = create_data_preparation(config_path, output_base_dir)
    if not data_preparation:
        logger.error("创建数据准备实例失败")
        return {}

    # 准备所有数据集
    all_experiment_data = data_preparation.prepare_all_datasets()

    # 输出准备结果统计
    total_anomalies = sum(
        len(dataset_data) for dataset_data in all_experiment_data.values()
    )
    logger.info(
        f"所有数据集准备完成，共 {len(all_experiment_data)} 个数据集，{total_anomalies} 条异常记录"
    )

    return all_experiment_data


if __name__ == "__main__":
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description="准备数据集的实验数据")
    parser.add_argument(
        "--config",
        type=str,
        default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml"),
        help="配置文件路径",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=os.path.join(
            Path(os.path.dirname(os.path.abspath(__file__))).parent.parent,
            "result",
            "telecom_processed_data",
        ),
        help="输出目录路径",
    )

    # 解析命令行参数
    args = parser.parse_args()

    # 显示参数信息
    logger.info(f"使用配置文件: {args.config}")
    logger.info(f"输出目录: {args.output}")

    # 确保输出目录存在
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 执行数据准备
    all_experiment_data = prepare_and_save_data(args.config, output_dir)

    # 输出最终统计信息
    total_datasets = len(all_experiment_data)
    total_anomalies = sum(
        len(dataset_data) for dataset_data in all_experiment_data.values()
    )

    logger.info("========== 数据准备完成 ==========")
    logger.info(f"共处理 {total_datasets} 个数据集，{total_anomalies} 条异常记录")
    for dataset_name, dataset_data in all_experiment_data.items():
        logger.info(f"  - 数据集 '{dataset_name}': {len(dataset_data)} 条异常记录")
